from langgraph.graph import StateGraph, END
from langgraph.checkpoint.sqlite import SqliteSaver

import sqlite3
from agents.onboarding_agent import OnboardingAgent
from typing import TypedDict, Annotated, Sequence
from config import REWARDS
from config import STAGES

class State(TypedDict):
    user_input: str
    agent_response: dict
    current_stage: str
    continue_conversation: bool

agent = OnboardingAgent()

def create_stage_node(stage_name):
    def stage_node(state: State):
        print(f"\n【{stage_name}阶段引导】")
        print(f"完成此阶段可得{REWARDS[stage_name]}")
        user_input = input("请输入回答或问题：")
        response = agent.process_input(user_input)
        return {
            **state,
            "user_input": user_input,
            "agent_response": response,
            "current_stage": stage_name,
            "continue_conversation": True
        }
    return stage_node

workflow = StateGraph(State)

for stage in ["注册", "开户", "入金", "首次交易"]:
    workflow.add_node(stage, create_stage_node(stage))

def should_continue(state: State):
    res = state["agent_response"]
    t = res.get("type")

    if t == "stage_progression":
        next_stage = STAGES[STAGES.index(state["current_stage"]) + 1]
        print(res["content"])
        return next_stage
    elif t == "onboarding_complete":
        print(res["content"])
        return END
    elif t == "rag":
        print(res["content"])
        print("\n来源：" + "\n".join(res["sources"]))
        print(res["follow_up"])
        return state["current_stage"]
    else:
        print(res["content"])
        return state["current_stage"]

for i, stage in enumerate(["注册", "开户", "入金", "首次交易"]):
    if i < 3:
        workflow.add_conditional_edges(
            stage,
            should_continue,
            {
                STAGES[i+1]: STAGES[i+1],
                END: END,
                stage: stage
            }
        )
    else:
        workflow.add_conditional_edges(
            stage,
            should_continue,
            {END: END, stage: stage}
        )

workflow.set_entry_point("注册")

# 创建 checkpointer 实例
sqlite_conn = sqlite3.connect("checkpoints.sqlite", check_same_thread=False)
checkpointer = SqliteSaver(sqlite_conn)
app = workflow.compile(checkpointer=checkpointer)
import os
os.environ["TOKENIZERS_PARALLELISM"] = "false"
# 编译流程
# 启用持久化 memory
print(f"[DEBUG] Checkpointer 类型：{type(checkpointer)}")